from .Strategy import Strategy
from openke.module.model import DistMult


class SUKE_Strategy(Strategy):

    def __init__(self, model=None, loss=None, batch_size=256, regul_rate=0.0, l3_regul_rate=0.0):
        super(SUKE_Strategy, self).__init__()
        self.model = model
        self.loss = loss
        self.batch_size = batch_size
        self.regul_rate = regul_rate
        self.l3_regul_rate = l3_regul_rate

    def _get_positive_score(self, score):
        positive_q_stru = score[0][:self.batch_size]
        positive_q_unce = score[1][:self.batch_size]
        positive_q_stru = positive_q_stru.view(-1, self.batch_size).permute(1, 0)
        positive_q_unce = positive_q_unce.view(-1, self.batch_size).permute(1, 0)
        return [positive_q_stru, positive_q_unce]

    def _get_negative_score(self, score):
        negative_q_stru = score[0][self.batch_size:]
        negative_q_unce = score[1][self.batch_size:]
        negative_q_stru = negative_q_stru.view(-1, self.batch_size).permute(1, 0)
        negative_q_unce = negative_q_unce.view(-1, self.batch_size).permute(1, 0)
        return [negative_q_stru, negative_q_unce]

    def forward(self, data):
        score = self.model(data)
        p_score = self._get_positive_score(score)
        n_score = self._get_negative_score(score)
        loss_res = self.loss(p_score, n_score)
        if self.regul_rate != 0:
            loss_res += self.regul_rate * self.model.regularization(data)
        if self.l3_regul_rate != 0:
            loss_res += self.l3_regul_rate * self.model.l3_regularization()
        return loss_res
